using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

using static Unity.Mathematics.math;
using quaternion = Unity.Mathematics.quaternion;

public class JobsFractal : MonoBehaviour
{
    [BurstCompile(FloatPrecision.Standard, FloatMode.Fast, CompileSynchronously = true)]
    private struct UpdateFractalLevelJob : IJobFor
    {
        public float spinAngleDelta;
        public float scale;

        [ReadOnly]
        public NativeArray<FractalPart> parents;
        public NativeArray<FractalPart> parts;
        [WriteOnly]
        public NativeArray<float3x4> matrices;
        
        public void Execute(int index)
        {
            var parent = parents[index / 5];
            var part = parts[index];
            part.spinAngle += spinAngleDelta;
            part.worldRotation = mul(parent.worldRotation, mul(part.rotation, quaternion.RotateY(part.spinAngle)));
            part.worldPosition = parent.worldPosition + mul(parent.worldRotation, 1.5f * scale * part.direction);
            parts[index] = part;
            float3x3 r = float3x3(part.worldRotation) * scale;
            matrices[index] = float3x4(r.c0, r.c1, r.c2, part.worldPosition);
        }
    }
    
    private struct FractalPart
    {
        public float3 direction, worldPosition;
        public quaternion rotation, worldRotation;
        public float spinAngle;
    }
    
    [SerializeField, Range(1, 8)]
    private int depth = 4;
    [SerializeField]
    private Mesh mesh;
    [SerializeField]
    private Material material;

    private static float3[] _directions = { up(), right(), left(), forward(), back() };
    private static quaternion[] _rotations =
    {
        quaternion.identity, quaternion.RotateZ(-0.5f * PI), quaternion.RotateZ(0.5f * PI),
        quaternion.RotateX(0.5f * PI), quaternion.RotateX(-0.5f * PI)
    };

    private static int _matricesId = Shader.PropertyToID("_Matrices");
    private static MaterialPropertyBlock propertyBlock;

    private NativeArray<FractalPart>[] _parts;
    private NativeArray<float3x4>[] _matrices;
    private ComputeBuffer[] _matricesBuffers;

    private void OnEnable()
    {
        _parts = new NativeArray<FractalPart>[depth];
        _matrices = new NativeArray<float3x4>[depth];
        _matricesBuffers = new ComputeBuffer[depth];
        int stride = 12 * 4;
        for (int i = 0, size = 1; i < _parts.Length; i++, size *= 5)
        {
            _parts[i] = new NativeArray<FractalPart>(size, Allocator.Persistent);
            _matrices[i] = new NativeArray<float3x4>(size, Allocator.Persistent);
            _matricesBuffers[i] = new ComputeBuffer(size, stride);
        }

        _parts[0][0] = CreatePart(0);
        for (int i = 1; i < _parts.Length; i++)
        {
            var levelParts = _parts[i];
            for (int j = 0; j < levelParts.Length; j += 5)
            {
                for (int k = 0; k < 5; k++)
                {
                    levelParts[j + k] = CreatePart(k);
                }
            }
        }

        if (propertyBlock == null)
            propertyBlock = new MaterialPropertyBlock();
    }

    private void OnDisable()
    {
        for (int i = 0; i < _matricesBuffers.Length; i++)
        {
            _matricesBuffers[i].Release();
            _parts[i].Dispose();
            _matrices[i].Dispose();
        }

        _parts = null;
        _matrices = null;
        _matricesBuffers = null;
    }

    private void OnValidate()
    {
        if (_parts != null && enabled)
        {
            OnDisable();
            OnEnable();
        }
    }

    private void Update()
    {
        var spinAngleDelta = 0.125f * PI * Time.deltaTime;
        var rootPart = _parts[0][0];
        rootPart.spinAngle += spinAngleDelta;
        var rootTransform = transform;
        rootPart.worldRotation =
            mul(rootTransform.rotation, mul(rootPart.rotation, quaternion.RotateY(rootPart.spinAngle)));
        rootPart.worldPosition = transform.position;
        _parts[0][0] = rootPart;

        float objectScale = rootTransform.lossyScale.x;
        var r = float3x3(rootPart.worldRotation) * objectScale;
        _matrices[0][0] = float3x4(r.c0, r.c1, r.c2, rootPart.worldPosition);

        var scale = objectScale;
        JobHandle jobHandler = default;
        for (int i = 1; i < _parts.Length; i++)
        {
            scale *= 0.5f;
            jobHandler = new UpdateFractalLevelJob()
            {
                spinAngleDelta = spinAngleDelta,
                scale = scale,
                parents = _parts[i - 1],
                parts = _parts[i],
                matrices = _matrices[i]
            }.ScheduleParallel(_parts[i].Length, 5, jobHandler);
        }
        jobHandler.Complete();

        var bounds = new Bounds(rootPart.worldPosition, 3f * objectScale * Vector3.one);
        for (int i = 0; i < _matricesBuffers.Length; i++)
        {
            var buffer = _matricesBuffers[i];
            buffer.SetData(_matrices[i]);
            propertyBlock.SetBuffer(_matricesId, buffer);
            Graphics.DrawMeshInstancedProcedural(mesh, 0, material, bounds, buffer.count, propertyBlock);
        }
    }

    private FractalPart CreatePart(int childIndex)
    {
        return new FractalPart
            {direction = _directions[childIndex], rotation = _rotations[childIndex]};
    }
}